"""
Week 5: 偏差-方差分解实现
Bias-Variance Decomposition Implementation
"""

import numpy as np
import matplotlib.pyplot as plt
from typing import Callable
from sklearn.preprocessing import PolynomialFeatures
import matplotlib
matplotlib.rcParams['font.sans-serif'] = ['SimHei']
matplotlib.rcParams['axes.unicode_minus'] = False

class BiasVarianceAnalyzer:
    """偏差-方差分析器"""
    
    def __init__(self):
        self.name = "Bias-Variance Analyzer"
    
    def true_function(self, x: np.ndarray) -> np.ndarray:
        """真实函数: f(x) = 1.5x - 0.5x² + 0.3x³"""
        return 1.5 * x - 0.5 * x**2 + 0.3 * x**3
    
    def generate_dataset(self, n_samples: int, noise_std: float = 0.3, 
                        x_range: tuple[float, float] = (-2, 2)) -> tuple[np.ndarray, np.ndarray]:
        """生成数据集"""
        x = np.random.uniform(x_range[0], x_range[1], n_samples)
        y_true = self.true_function(x)
        noise = np.random.normal(0, noise_std, n_samples)
        y = y_true + noise
        return x.reshape(-1, 1), y
    
    def polynomial_model(self, X: np.ndarray, y: np.ndarray, degree: int) -> Callable:
        """多项式模型拟合"""
        # 创建多项式特征
        poly_features = PolynomialFeatures(degree=degree, include_bias=True)
        X_poly = poly_features.fit_transform(X)
        
        # 最小二乘拟合
        try:
            theta = np.linalg.inv(X_poly.T @ X_poly) @ X_poly.T @ y
        except np.linalg.LinAlgError:
            theta = np.linalg.pinv(X_poly) @ y
        
        def predict(x_new):
            x_new = x_new.reshape(-1, 1)
            X_new_poly = poly_features.transform(x_new)
            return X_new_poly @ theta
        
        return predict
    
    def bias_variance_decomposition(self, degrees: list[int], n_experiments: int = 100, 
                                  n_samples: int = 50, noise_std: float = 0.3):
        """执行偏差-方差分解实验"""
        x_test = np.linspace(-2, 2, 100)
        y_true = self.true_function(x_test)
        
        results = {}
        
        for degree in degrees:
            print(f"分析多项式度数 {degree}...")
            
            predictions = []
            
            # 进行多次实验
            for exp in range(n_experiments):
                # 生成训练数据
                X_train, y_train = self.generate_dataset(n_samples, noise_std)
                
                # 训练模型
                model = self.polynomial_model(X_train, y_train, degree)
                
                # 在测试点预测
                y_pred = model(x_test)
                predictions.append(y_pred)
            
            # 转换为数组
            predictions = np.array(predictions)  # shape: (n_experiments, n_test_points)
            
            # 计算偏差、方差和总误差
            mean_prediction = np.mean(predictions, axis=0)
            
            # 偏差²
            bias_squared = np.mean((mean_prediction - y_true) ** 2)
            
            # 方差
            variance = np.mean(np.var(predictions, axis=0))
            
            # 噪声 (理论值)
            noise = noise_std ** 2
            
            # 总误差 (经验估计)
            total_error = np.mean(np.mean((predictions - y_true) ** 2, axis=1))
            
            results[degree] = {
                'bias_squared': bias_squared,
                'variance': variance,
                'noise': noise,
                'total_error': total_error,
                'predictions': predictions,
                'mean_prediction': mean_prediction
            }
            
            print(f"  偏差²: {bias_squared:.4f}")
            print(f"  方差: {variance:.4f}")
            print(f"  噪声: {noise:.4f}")
            print(f"  总误差: {total_error:.4f}")
            print(f"  理论总误差: {bias_squared + variance + noise:.4f}\n")
        
        return results, x_test, y_true
    
    def visualize_bias_variance(self, results: dict, x_test: np.ndarray, y_true: np.ndarray):
        """可视化偏差-方差分解结果"""
        degrees = list(results.keys())
        
        fig, axes = plt.subplots(2, 3, figsize=(18, 12))
        
        # 1. 不同复杂度模型的拟合效果
        ax1 = axes[0, 0]
        ax1.plot(x_test, y_true, 'k-', linewidth=3, label='真实函数')
        
        colors = ['blue', 'red', 'green', 'orange', 'purple']
        for i, degree in enumerate(degrees[:5]):
            result = results[degree]
            mean_pred = result['mean_prediction']
            predictions = result['predictions']
            
            # 绘制多个预测的样本
            for j in range(min(10, predictions.shape[0])):
                ax1.plot(x_test, predictions[j], color=colors[i], alpha=0.1, linewidth=1)
            
            # 绘制平均预测
            ax1.plot(x_test, mean_pred, color=colors[i], linewidth=2, 
                    label=f'度数{degree} (平均)')
        
        ax1.set_xlabel('x')
        ax1.set_ylabel('y')
        ax1.set_title('不同复杂度模型的预测')
        ax1.legend()
        ax1.grid(True, alpha=0.3)
        
        # 2. 偏差²随模型复杂度变化
        ax2 = axes[0, 1]
        bias_values = [results[d]['bias_squared'] for d in degrees]
        ax2.plot(degrees, bias_values, 'ro-', linewidth=2, markersize=8)
        ax2.set_xlabel('多项式度数')
        ax2.set_ylabel('偏差²')
        ax2.set_title('偏差² vs 模型复杂度')
        ax2.grid(True, alpha=0.3)
        
        # 3. 方差随模型复杂度变化
        ax3 = axes[0, 2]
        variance_values = [results[d]['variance'] for d in degrees]
        ax3.plot(degrees, variance_values, 'bo-', linewidth=2, markersize=8)
        ax3.set_xlabel('多项式度数')
        ax3.set_ylabel('方差')
        ax3.set_title('方差 vs 模型复杂度')
        ax3.grid(True, alpha=0.3)
        
        # 4. 偏差-方差权衡
        ax4 = axes[1, 0]
        noise_values = [results[d]['noise'] for d in degrees]
        total_error_values = [results[d]['total_error'] for d in degrees]
        
        ax4.plot(degrees, bias_values, 'r-', linewidth=2, marker='o', label='偏差²')
        ax4.plot(degrees, variance_values, 'b-', linewidth=2, marker='s', label='方差')
        ax4.plot(degrees, noise_values, 'g--', linewidth=2, label='噪声')
        ax4.plot(degrees, total_error_values, 'k-', linewidth=3, marker='^', label='总误差')
        
        ax4.set_xlabel('多项式度数')
        ax4.set_ylabel('误差')
        ax4.set_title('偏差-方差权衡')
        ax4.legend()
        ax4.grid(True, alpha=0.3)
        
        # 5. 误差分解饼图 (选择中等复杂度)
        ax5 = axes[1, 1]
        mid_degree = degrees[len(degrees)//2]
        result = results[mid_degree]
        
        sizes = [result['bias_squared'], result['variance'], result['noise']]
        labels = ['偏差²', '方差', '噪声']
        colors_pie = ['red', 'blue', 'green']
        
        ax5.pie(sizes, labels=labels, colors=colors_pie, autopct='%1.1f%%', startangle=90)
        ax5.set_title(f'误差分解 (度数={mid_degree})')
        
        # 6. 预测区间可视化
        ax6 = axes[1, 2]
        ax6.plot(x_test, y_true, 'k-', linewidth=3, label='真实函数')
        
        # 选择一个中等复杂度的模型
        result = results[mid_degree]
        mean_pred = result['mean_prediction']
        predictions = result['predictions']
        
        # 计算预测区间
        pred_std = np.std(predictions, axis=0)
        ax6.fill_between(x_test, mean_pred - 2*pred_std, mean_pred + 2*pred_std, 
                        alpha=0.3, color='blue', label='95%预测区间')
        ax6.plot(x_test, mean_pred, 'b-', linewidth=2, label=f'平均预测 (度数={mid_degree})')
        
        ax6.set_xlabel('x')
        ax6.set_ylabel('y')
        ax6.set_title('预测不确定性')
        ax6.legend()
        ax6.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.show()

class NoiseAnalyzer:
    """噪声分析器"""
    
    def __init__(self):
        self.name = "Noise Analyzer"
    
    def analyze_noise_impact(self, noise_levels: list[float], model_degrees: list[int]):
        """分析噪声对不同复杂度模型的影响"""
        analyzer = BiasVarianceAnalyzer()
        
        results = {}
        
        for noise_std in noise_levels:
            print(f"分析噪声水平: σ = {noise_std}")
            
            bias_variance_results, x_test, y_true = analyzer.bias_variance_decomposition(
                degrees=model_degrees, 
                n_experiments=50, 
                n_samples=30, 
                noise_std=noise_std
            )
            
            results[noise_std] = bias_variance_results
        
        self.visualize_noise_impact(results, noise_levels, model_degrees)
        
        return results
    
    def visualize_noise_impact(self, results: dict, noise_levels: list[float], 
                             model_degrees: list[int]):
        """可视化噪声影响"""
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        
        # 1. 偏差随噪声变化
        ax1 = axes[0, 0]
        for degree in model_degrees:
            bias_values = [results[noise][degree]['bias_squared'] for noise in noise_levels]
            ax1.plot(noise_levels, bias_values, 'o-', label=f'度数 {degree}', linewidth=2)
        
        ax1.set_xlabel('噪声标准差')
        ax1.set_ylabel('偏差²')
        ax1.set_title('偏差² vs 噪声水平')
        ax1.legend()
        ax1.grid(True, alpha=0.3)
        
        # 2. 方差随噪声变化
        ax2 = axes[0, 1]
        for degree in model_degrees:
            variance_values = [results[noise][degree]['variance'] for noise in noise_levels]
            ax2.plot(noise_levels, variance_values, 's-', label=f'度数 {degree}', linewidth=2)
        
        ax2.set_xlabel('噪声标准差')
        ax2.set_ylabel('方差')
        ax2.set_title('方差 vs 噪声水平')
        ax2.legend()
        ax2.grid(True, alpha=0.3)
        
        # 3. 总误差随噪声变化
        ax3 = axes[1, 0]
        for degree in model_degrees:
            total_error_values = [results[noise][degree]['total_error'] for noise in noise_levels]
            ax3.plot(noise_levels, total_error_values, '^-', label=f'度数 {degree}', linewidth=2)
        
        ax3.set_xlabel('噪声标准差')
        ax3.set_ylabel('总误差')
        ax3.set_title('总误差 vs 噪声水平')
        ax3.legend()
        ax3.grid(True, alpha=0.3)
        
        # 4. 信噪比分析
        ax4 = axes[1, 1]
        # 计算信号功率 (真实函数的方差)
        x_test = np.linspace(-2, 2, 100)
        analyzer = BiasVarianceAnalyzer()
        y_true = analyzer.true_function(x_test)
        signal_power = np.var(y_true)
        
        snr_values = [signal_power / (noise**2) for noise in noise_levels]
        snr_db = [10 * np.log10(snr) for snr in snr_values]
        
        # 选择一个中等复杂度模型
        mid_degree = model_degrees[len(model_degrees)//2]
        total_errors = [results[noise][mid_degree]['total_error'] for noise in noise_levels]
        
        ax4.plot(snr_db, total_errors, 'ro-', linewidth=2, markersize=8)
        ax4.set_xlabel('信噪比 (dB)')
        ax4.set_ylabel('总误差')
        ax4.set_title(f'误差 vs 信噪比 (度数={mid_degree})')
        ax4.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.show()

def demonstrate_bias_variance():
    """演示偏差-方差分解"""
    print("=== 偏差-方差分解演示 ===\n")
    
    analyzer = BiasVarianceAnalyzer()
    
    # 分析不同复杂度的模型
    degrees = [1, 2, 3, 5, 8, 12]
    
    print("执行偏差-方差分解实验...")
    results, x_test, y_true = analyzer.bias_variance_decomposition(
        degrees=degrees, 
        n_experiments=100, 
        n_samples=50, 
        noise_std=0.3
    )
    
    # 可视化结果
    analyzer.visualize_bias_variance(results, x_test, y_true)
    
    # 找到最优复杂度
    total_errors = [results[d]['total_error'] for d in degrees]
    optimal_degree = degrees[np.argmin(total_errors)]
    
    print(f"最优模型复杂度: 度数 {optimal_degree}")
    print(f"最小总误差: {min(total_errors):.4f}")

def demonstrate_noise_analysis():
    """演示噪声分析"""
    print("\n=== 噪声影响分析 ===\n")
    
    noise_analyzer = NoiseAnalyzer()
    
    # 不同噪声水平
    noise_levels = [0.1, 0.2, 0.3, 0.5, 0.8]
    model_degrees = [1, 3, 5, 8]
    
    print("分析噪声对不同复杂度模型的影响...")
    results = noise_analyzer.analyze_noise_impact(noise_levels, model_degrees)
    
    # 分析结果
    print("\n噪声影响总结:")
    print("- 噪声主要影响方差，对偏差影响较小")
    print("- 复杂模型对噪声更敏感")
    print("- 高噪声环境下应选择简单模型")

if __name__ == "__main__":
    demonstrate_bias_variance()
    demonstrate_noise_analysis()